import json

from main import chain
import time


result_dict = {}
total_time = 0


def test():

    global total_time
    global result_dict

    begin_t = time.time()
    answer = chain.invoke(input={'sector': '手表行业'})
    end_t = time.time()

    total_time += end_t - begin_t

    cnt = answer.content

    for line in cnt.split():
        if line.count('-') == 1:
            key, val = line.split('-')
            values = sorted(val.strip('。').split('、'))
            if key in result_dict:
                if values not in result_dict[key]:
                    result_dict[key].append(values)
            else:
                result_dict[key] = [values]


if __name__ == '__main__':
    test_times = 100
    for i in range(test_times):
        print(f'\rtest: {i+1}', end='')
        test()

    print()
    for key, value in result_dict.items():
        print(f'{key}: { (len(value) - 1) / test_times * 100:.2f}%')

    print(f'avg time: {total_time / test_times:.2f}')

    with open('./test.json', 'w', encoding='utf-8') as f:
        json.dump(result_dict, f, ensure_ascii=False, indent=4)


